use num_traits::Num;
use crate::autograd::function::{ForwardArgs, Function};
use crate::autograd::function_ctx::FunctionCtx;
use crate::Tensor;

pub struct AddBackward<T> {
    ctx: FunctionCtx<T>,
}

impl<T> Function<T> for AddBackward<T>
    where
        T: Copy + Num + 'static,
{
    fn new(ctx: FunctionCtx<T>) -> Self {
        Self { ctx }
    }

    fn ctx(&self) -> &FunctionCtx<T> {
        &self.ctx
    }

    fn forward(_ctx: &mut FunctionCtx<T>, args: ForwardArgs<T>) -> Tensor<T> {
        if let ForwardArgs::TensorTensor(lhs, rhs) = args {
            return lhs.zip(rhs, |a, b| a + b)
        }
        unreachable!()
    }

    fn backward(&self, grad_output: Tensor<T>) -> Vec<Option<Tensor<T>>> {
        [grad_output.clone().into(), grad_output.into()].into()
    }
}
